{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to integrate Langchain4j LLM streaming in Langgraph4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "var userHomeDir = System.getProperty(\"user.home\");\n",
    "var localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "var langchain4jVersion = \"1.0.1\";\n",
    "var langchain4jbeta = \"1.0.1-beta6\";\n",
    "var langgraph4jVersion = \"1.6-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove installed package from Jupiter cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos\n",
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "// %dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-ollama:\\{langchain4jbeta}\n",
    "\n",
    "\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize Logger**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    java.util.logging.LogManager.getLogManager().readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"llm-streaming\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to use StreamingChatGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "StreamingOutput{chunk=Here} \n",
      "StreamingOutput{chunk='s} \n",
      "StreamingOutput{chunk= one} \n",
      "StreamingOutput{chunk=:\n",
      "\n",
      "} \n",
      "StreamingOutput{chunk=What} \n",
      "StreamingOutput{chunk= do} \n",
      "StreamingOutput{chunk= you} \n",
      "StreamingOutput{chunk= call} \n",
      "StreamingOutput{chunk= a} \n",
      "StreamingOutput{chunk= fake} \n",
      "StreamingOutput{chunk= nood} \n",
      "StreamingOutput{chunk=le} \n",
      "StreamingOutput{chunk=?\n",
      "\n",
      "} \n",
      "StreamingOutput{chunk=An} \n",
      "StreamingOutput{chunk= imp} \n",
      "StreamingOutput{chunk=asta} \n",
      "StreamingOutput{chunk=!\n",
      "\n",
      "} \n",
      "StreamingOutput{chunk=Hope} \n",
      "StreamingOutput{chunk= that} \n",
      "StreamingOutput{chunk= made} \n",
      "StreamingOutput{chunk= you} \n",
      "StreamingOutput{chunk= laugh} \n",
      "StreamingOutput{chunk=!} \n",
      "StreamingOutput{chunk= Do} \n",
      "StreamingOutput{chunk= you} \n",
      "StreamingOutput{chunk= want} \n",
      "StreamingOutput{chunk= to} \n",
      "StreamingOutput{chunk= hear} \n",
      "StreamingOutput{chunk= another} \n",
      "StreamingOutput{chunk= one} \n",
      "StreamingOutput{chunk=?} \n",
      "RESULT: {content=Here's one:\n",
      "\n",
      "What do you call a fake noodle?\n",
      "\n",
      "An impasta!\n",
      "\n",
      "Hope that made you laugh! Do you want to hear another one?} \n"
     ]
    }
   ],
   "source": [
    "import dev.langchain4j.model.StreamingResponseHandler;\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.model.output.Response;\n",
    "import org.bsc.langgraph4j.langchain4j.generators.StreamingChatGenerator;\n",
    "import org.bsc.langgraph4j.state.AgentState;\n",
    "import org.bsc.langgraph4j.streaming.StreamingOutput;\n",
    "import dev.langchain4j.model.ollama.OllamaStreamingChatModel;\n",
    "import dev.langchain4j.model.chat.request.ChatRequest;\n",
    "\n",
    "var generator = StreamingChatGenerator.<AgentState>builder()\n",
    "                        .mapResult( r -> Map.of( \"content\", r.aiMessage().text() ) )\n",
    "                        .build();\n",
    "\n",
    "var model = OllamaStreamingChatModel.builder()\n",
    "    .baseUrl( \"http://localhost:11434\" )\n",
    "    .temperature(0.0)\n",
    "    .logRequests(true)\n",
    "    .logResponses(true)\n",
    "    .modelName(\"llama3.1:latest\")\n",
    "    .build();\n",
    "\n",
    "var request = ChatRequest.builder()\n",
    "        .messages( UserMessage.from(\"Tell me a joke\") )\n",
    "        .build();\n",
    "model.chat(request, generator.handler() );\n",
    "\n",
    "for( var r : generator ) {\n",
    "    log.info( \"{}\", r);\n",
    "}\n",
    "  \n",
    "log.info( \"RESULT: {}\", generator.resultValue().orElse(null) );\n",
    "  \n",
    "//Thread.sleep( 1000 );"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use StreamingChatGenerator in Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Serializers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.SystemMessage;\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.data.message.ToolExecutionResultMessage;\n",
    "import dev.langchain4j.agent.tool.ToolExecutionRequest;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.LC4jStateSerializer;\n",
    "import org.bsc.langgraph4j.state.AgentStateFactory;\n",
    "import org.bsc.langgraph4j.prebuilt.MessagesState;\n",
    "\n",
    "var stateSerializer = new LC4jStateSerializer<MessagesState<ChatMessage>>( MessagesState::new );\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the tools\n",
    "\n",
    "Using [langchain4j], We will first define the tools we want to use. For this simple example, we will\n",
    "use create a placeholder search engine. However, it is really easy to create\n",
    "your own tools - see documentation\n",
    "[here][tools] on how to do\n",
    "that.\n",
    "\n",
    "[langchain4j]: https://docs.langchain4j.dev\n",
    "[tools]: https://docs.langchain4j.dev/tutorials/tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.agent.tool.P;\n",
    "import dev.langchain4j.agent.tool.Tool;\n",
    "\n",
    "import java.util.Optional;\n",
    "\n",
    "import static java.lang.String.format;\n",
    "\n",
    "public class SearchTool {\n",
    "\n",
    "    @Tool(\"get weather realtime information.\")\n",
    "    String execQuery(@P(\"cit.\") String city) {\n",
    "\n",
    "        // This is a placeholder for the actual implementation\n",
    "        return \"Cold, with a low of 13 degrees\";\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import org.bsc.langgraph4j.prebuilt.MessagesStateGraph;\n",
    "import org.bsc.langgraph4j.action.EdgeAction;\n",
    "import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;\n",
    "import org.bsc.langgraph4j.action.NodeAction;\n",
    "import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;\n",
    "import dev.langchain4j.service.tool.DefaultToolExecutor;\n",
    "import org.bsc.langgraph4j.langchain4j.tool.LC4jToolService;\n",
    "import dev.langchain4j.agent.tool.ToolSpecification;\n",
    "import dev.langchain4j.model.chat.request.ChatRequestParameters;\n",
    "\n",
    "// setup streaming model\n",
    "var model = OllamaStreamingChatModel.builder()\n",
    "    .baseUrl( \"http://localhost:11434\" )\n",
    "    .temperature(0.0)\n",
    "    .logRequests(true)\n",
    "    .logResponses(true)\n",
    "    .modelName(\"llama3.1:latest\")\n",
    "    .build();\n",
    "\n",
    "// setup tools \n",
    "var tools = LC4jToolService.builder()\n",
    "              .toolsFromObject( new SearchTool() )\n",
    "              .build();\n",
    "\n",
    "NodeAction<MessagesState<ChatMessage>> callModel = state -> {\n",
    "        log.info( \"CallModel\" );\n",
    "\n",
    "        var generator = StreamingChatGenerator.<MessagesState<ChatMessage>>builder()\n",
    "                .mapResult( response -> Map.of(\"messages\", response.aiMessage()) )\n",
    "                .startingNode(\"agent\")\n",
    "                .startingState(state)\n",
    "                .build();\n",
    "\n",
    "        var parameters = ChatRequestParameters.builder()\n",
    "                .toolSpecifications(tools.toolSpecifications())\n",
    "                .build();  \n",
    "\n",
    "        var request = ChatRequest.builder()\n",
    "                .parameters(parameters)\n",
    "                .messages( state.messages() )\n",
    "                .build();\n",
    "        \n",
    "        model.chat( request, generator.handler() );\n",
    "\n",
    "        return Map.of(\"_streaming_messages\", generator);\n",
    "};\n",
    "            \n",
    "// Route Message\n",
    "EdgeAction<MessagesState<ChatMessage>> routeMessage = state -> {\n",
    "\n",
    "        var lastMessage = state.lastMessage()\n",
    "            .orElseThrow(() -> (new IllegalStateException(\"last message not found!\")));\n",
    "\n",
    "        log.info(\"routeMessage:\\n{}\", lastMessage );\n",
    "        \n",
    "        if (lastMessage instanceof AiMessage message) {\n",
    "                // If tools should be called\n",
    "                if (message.hasToolExecutionRequests()) { \n",
    "                        return \"next\";\n",
    "                }\n",
    "        }\n",
    "\n",
    "        // If no tools are called, we can finish (respond to the user)\n",
    "        return \"exit\";\n",
    "};\n",
    "            \n",
    "// Invoke Tool\n",
    "NodeAction<MessagesState<ChatMessage>> invokeTool = state -> {\n",
    "\n",
    "    var lastMessage = state.lastMessage()\n",
    "            .orElseThrow(() -> (new IllegalStateException(\"last message not found!\")));\n",
    "\n",
    "    log.info(\"invokeTool:\\n{}\", lastMessage );\n",
    "\n",
    "    if (lastMessage instanceof AiMessage lastAiMessage) {\n",
    "\n",
    "        var result = tools.execute(lastAiMessage.toolExecutionRequests(), null)\n",
    "                .orElseThrow(() -> (new IllegalStateException(\"no tool found!\")));\n",
    "\n",
    "        return Map.of(\"messages\", result);\n",
    "\n",
    "    }\n",
    "\n",
    "    throw new IllegalStateException(\"invalid last message\");\n",
    "};\n",
    "            \n",
    "// Define Graph\n",
    "var workflow = new MessagesStateGraph<ChatMessage>(stateSerializer)\n",
    "        .addNode(\"agent\", node_async(callModel))\n",
    "        .addNode(\"tools\", node_async(invokeTool))\n",
    "        .addEdge(START, \"agent\")\n",
    "        .addConditionalEdges(\"agent\",\n",
    "                edge_async(routeMessage),\n",
    "                Map.of(\"next\", \"tools\", \"exit\", END))\n",
    "        .addEdge(\"tools\", \"agent\");\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "CallModel \n",
      "routeMessage:\n",
      "AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] } \n",
      "ToolExecutionRequest id is null! \n",
      "NodeOutput{node=__START__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"what is the whether in Napoli?\" }] }]}} \n",
      "ToolExecutionRequest id is null! \n",
      "invokeTool:\n",
      "AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] } \n",
      "execute: execQuery \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "NodeOutput{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"what is the whether in Napoli?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] }]}} \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "CallModel \n",
      "NodeOutput{node=tools, state={messages=[UserMessage { name = null contents = [TextContent { text = \"what is the whether in Napoli?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] }, ToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }]}} \n",
      "StreamingOutput{node=agent, chunk=The } \n",
      "StreamingOutput{node=agent, chunk= current } \n",
      "StreamingOutput{node=agent, chunk= weather } \n",
      "StreamingOutput{node=agent, chunk= in } \n",
      "StreamingOutput{node=agent, chunk= Napoli } \n",
      "StreamingOutput{node=agent, chunk= is } \n",
      "StreamingOutput{node=agent, chunk= cold } \n",
      "StreamingOutput{node=agent, chunk=, } \n",
      "StreamingOutput{node=agent, chunk= with } \n",
      "StreamingOutput{node=agent, chunk= a } \n",
      "StreamingOutput{node=agent, chunk= low } \n",
      "StreamingOutput{node=agent, chunk= temperature } \n",
      "StreamingOutput{node=agent, chunk= of } \n",
      "StreamingOutput{node=agent, chunk=  } \n",
      "StreamingOutput{node=agent, chunk=13 } \n",
      "StreamingOutput{node=agent, chunk= degrees } \n",
      "routeMessage:\n",
      "AiMessage { text = \"The current weather in Napoli is cold, with a low temperature of 13 degrees.\" toolExecutionRequests = [] } \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "StreamingOutput{node=agent, chunk=. } \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "NodeOutput{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"what is the whether in Napoli?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] }, ToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in Napoli is cold, with a low temperature of 13 degrees.\" toolExecutionRequests = [] }]}} \n",
      "NodeOutput{node=__END__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"what is the whether in Napoli?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"city\" : \"Napoli\"\n",
      "}\" }] }, ToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in Napoli is cold, with a low temperature of 13 degrees.\" toolExecutionRequests = [] }]}} \n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.streaming.StreamingOutput;\n",
    "\n",
    "var app = workflow.compile();\n",
    "\n",
    "for( var out : app.stream( Map.of( \"messages\", UserMessage.from( \"what is the whether in Napoli?\")) ) ) {\n",
    "  if( out instanceof StreamingOutput streaming ) {\n",
    "    log.info( \"StreamingOutput{node={}, chunk={} }\", streaming.node(), streaming.chunk() );\n",
    "  }\n",
    "  else {\n",
    "    log.info( \"{}\", out );\n",
    "  }\n",
    "}\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "22.0.2+9-70"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
